#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import pymongo
import click
from tqdm import tqdm

DEFAULT_MONGO_URI = "mongodb://mongo-wepd.wepd:27017/wepd"
DEFAULT_REF_COLLECTION = "ref"


@click.command()
@click.option("--mongo_uri", "-m", help="MongoDB URI", required=True, type=str, default=DEFAULT_MONGO_URI,
              show_default=True)
@click.option("--ref", "-r", help="Reference collection", required=True, type=str, default=DEFAULT_REF_COLLECTION,
              show_default=True)
@click.argument("ref_path", type=click.Path(exists=True))
def wepd_ref2mongo(mongo_uri: str, ref: str, ref_path: str):
    """
    import reference fasta into MongoDB(single chromosome)
    :param mongo_uri: MongoDB connection URI
    :param ref: MongoDB reference collection
    :param ref_path: reference fasta path(single chromosome)
    :return: none
    """
    client = pymongo.MongoClient(mongo_uri)
    db = client.get_default_database()
    collection = db[ref]

    current_ref = ""
    current_start = 1
    infile = open(ref_path, "r")
    title = infile.readline().rstrip()
    chrom = title[1:]
    for line in tqdm(infile, desc="line"):
        line = line.rstrip()
        current_ref += line
        if len(current_ref) >= 1500:
            collection.insert_one({"chrom_id": chrom, "start": current_start, "end": current_start + 1499,
                                   "sequence": current_ref[:1500].upper()})
            current_start += 1500
            current_ref = current_ref[1500:]
    if len(current_ref) > 0:
        collection.insert_one({"chrom_id": chrom, "start": current_start, "end": current_start + len(current_ref) - 1,
                               "sequence": current_ref.upper()})
    infile.close()
    client.close()


if __name__ == '__main__':
    wepd_ref2mongo()
